#!/usr/bin/env python3
"""
yolo_crop.py

根据 YOLO 格式的 label 文件（每行: class x_center y_center w h，均为归一化值）
从 images 目录裁剪出每个标注的目标，保存到 out_dir。

输出命名: <原图文件名>_<num><原后缀>
例如: dog.jpg -> dog_0.jpg, dog_1.jpg

用法:
    python yolo_crop.py --images /path/to/images --labels /path/to/labels --out /path/to/output

可选参数:
    --min-size N    最小保存宽/高阈值（默认1）
    --exts .jpg .png 搜索 images 的扩展名列表（按优先级）
"""

import argparse
import os
import glob
from PIL import Image

COMMON_EXTS = [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"]

def find_image_by_stem(images_dir, stem, exts=None):
    """在 images_dir 中按优先级查找与 stem 匹配的图片文件（不递归）。返回第一个匹配路径或 None。"""
    exts = exts or COMMON_EXTS
    for e in exts:
        p = os.path.join(images_dir, stem + e)
        if os.path.isfile(p):
            return p
    # 也尝试大小写不同的扩展或实际文件名不同的情况
    # 列出 images_dir 中所有以 stem 开头的文件，做后缀比较（稳健）
    for fname in os.listdir(images_dir):
        if not os.path.isfile(os.path.join(images_dir, fname)):
            continue
        name, ext = os.path.splitext(fname)
        if name == stem:
            return os.path.join(images_dir, fname)
    return None

def parse_label_line(line):
    """
    解析一行 YOLO label，返回 (cls, x_center, y_center, w, h) 或 None（解析失败）
    允许行首尾空格。
    """
    parts = line.strip().split()
    if len(parts) < 5:
        return None
    try:
        cls = parts[0]
        xc = float(parts[1])
        yc = float(parts[2])
        w = float(parts[3])
        h = float(parts[4])
        return cls, xc, yc, w, h
    except Exception:
        return None

def crop_and_save(img_path, bbox_norm, out_path):
    """
    bbox_norm: (xc, yc, w, h) in normalized coords [0,1]
    out_path: full path to save
    """
    with Image.open(img_path) as im:
        W, H = im.size
        xc, yc, w, h = bbox_norm
        # convert to pixel coords
        x1 = (xc - w/2) * W
        y1 = (yc - h/2) * H
        x2 = (xc + w/2) * W
        y2 = (yc + h/2) * H
        # clamp to image bounds and to integer pixel indices
        x1 = int(max(0, round(x1)))
        y1 = int(max(0, round(y1)))
        x2 = int(min(W, round(x2)))
        y2 = int(min(H, round(y2)))
        # ensure non-empty box (PIL crop requires box where x2>x1 and y2>y1)
        if x2 <= x1 or y2 <= y1:
            # fallback: try expanding to at least 1 pixel within bounds
            if x1 < W:
                x2 = min(W, x1 + 1)
            else:
                x1 = max(0, W - 1); x2 = W
            if y1 < H:
                y2 = min(H, y1 + 1)
            else:
                y1 = max(0, H - 1); y2 = H
        cropped = im.crop((x1, y1, x2, y2))
        # preserve format by saving with extension in out_path
        cropped.save(out_path)

def collect_label_files(labels_dir):
    """返回 labels_dir 下所有以 .txt 结尾的文件的完整路径（不递归），按文件名顺序排序。"""
    files = []
    for fname in os.listdir(labels_dir):
        fpath = os.path.join(labels_dir, fname)
        if os.path.isfile(fpath) and fname.lower().endswith(".txt"):
            files.append(fpath)
    files.sort()  # 保持稳定的原始排序
    return files

def main(args):
    images_dir = args.images
    labels_dir = args.labels
    out_dir = args.out
    exts = args.exts if args.exts else COMMON_EXTS
    min_size = int(args.min_size)

    if not os.path.isdir(images_dir):
        raise SystemExit(f"images 目录不存在: {images_dir}")
    if not os.path.isdir(labels_dir):
        raise SystemExit(f"labels 目录不存在: {labels_dir}")
    os.makedirs(out_dir, exist_ok=True)

    label_files = collect_label_files(labels_dir)
    if not label_files:
        print("labels 目录中没有 .txt 文件。")
        return

    total_crops = 0
    skipped_no_image = 0
    for label_path in label_files:
        label_name = os.path.basename(label_path)
        stem, _ = os.path.splitext(label_name)
        img_path = find_image_by_stem(images_dir, stem, exts)
        if img_path is None:
            skipped_no_image += 1
            print(f"[跳过] 未找到对应图像: {label_name} -> images 中没有找到与 '{stem}' 匹配的文件。")
            continue

        # 读取所有 label 行（先收集再处理）
        with open(label_path, "r", encoding="utf-8") as f:
            lines = [ln.rstrip("\n") for ln in f]

        # 如果没有标注（空文件或仅空行），跳过
        nonempty_lines = [ln for ln in lines if ln.strip()]
        if not nonempty_lines:
            # 可选：如果需要生成整张图像作为单独裁剪，可以在这里处理。目前按要求跳过。
            print(f"[跳过] label 为空: {label_name}")
            continue

        # 处理每一行
        crop_idx = 0
        for ln in nonempty_lines:
            parsed = parse_label_line(ln)
            if parsed is None:
                # 非法行，跳过
                print(f"[警告] 无法解析行（跳过）: {label_path}: '{ln}'")
                continue
            cls, xc, yc, w, h = parsed
            # 计算 pixel bbox 在图像中大小以判断是否低于 min_size
            # 为此需要图像尺寸 — 读取一次尺寸（Pillow 可用 Image.open().size）
            with Image.open(img_path) as tmp_im:
                W, H = tmp_im.size
            px_w = int(round(w * W))
            px_h = int(round(h * H))
            if px_w < min_size or px_h < min_size:
                print(f"[略过小框] {label_name} 行: '{ln}' -> 宽高({px_w},{px_h}) < min_size({min_size})")
                continue

            # 构造输出文件名：<原名>_<num><原后缀>
            _, img_ext = os.path.splitext(img_path)
            safe_ext = img_ext if img_ext else ".jpg"
            out_fname = f"{stem}_{crop_idx}{safe_ext}"
            out_path = os.path.join(out_dir, out_fname)

            try:
                crop_and_save(img_path, (xc, yc, w, h), out_path)
                total_crops += 1
                crop_idx += 1
            except Exception as e:
                print(f"[错误] 裁剪/保存失败: {img_path} 行 '{ln}' -> {e}")

        if crop_idx == 0:
            # 如果一个 label 文件有标注行但全部被过滤（例如都小于 min_size）
            print(f"[信息] {label_name} 没有生成裁剪（可能都被过滤）")

    print("完成。")
    print(f"总生成裁剪数: {total_crops}")
    if skipped_no_image:
        print(f"跳过的 label 文件（未找到对应图像）: {skipped_no_image}")

if __name__ == "__main__":
    p = argparse.ArgumentParser(description="根据 YOLO label 从 images 中裁剪目标并保存")
    p.add_argument("--images", required=True, help="images 目录（YOLO 图像路径）")
    p.add_argument("--labels", required=True, help="labels 目录（YOLO label txt）")
    p.add_argument("--out", required=True, help="裁剪图像输出目录")
    p.add_argument("--min-size", default=1, help="最小裁剪宽或高（像素），低于则跳过（默认1）")
    p.add_argument("--exts", nargs="*", help="按优先级查找图像的扩展名（如 .jpg .png），默认常见格式")
    args = p.parse_args()
    main(args)
